import matplotlib.pyplot as plt

x_data = [[1, 2, 1, 1], [2, 1, 3, 2], [3, 1, 3, 4], [4, 1, 5, 5],
          [1, 7, 5, 5], [1, 2, 5, 6], [1, 6, 6, 6], [1, 7, 7, 7]]
y_data = [[0, 0, 1], [0, 0, 1], [0, 0, 1], [0, 1, 0],
          [0, 1, 0], [0, 1, 0], [1, 0, 0], [1, 0, 0]]
# 1.使用pytorch，完成多分类处理
import torch

# (1)数据处理
# ①将上面数据导入
x = torch.Tensor(x_data)
y = torch.Tensor(y_data)
x.shape
print(x.shape)
print(y.shape)
# ②创建多分类模型
model = torch.nn.Linear(in_features=x.shape[1], out_features=y.shape[1])
loss_fn = torch.nn.CrossEntropyLoss()
# ③实现梯度下降
op = torch.optim.Adam(model.parameters(), lr=0.01)
# ④底层实现多分类函数
# ⑤每200次打印代价函数一次
loss_list = []
for i in range(20000):
    op.zero_grad()
    h = model(x)
    loss = loss_fn(h, y)
    loss.backward()
    op.step()
    loss_list.append(loss)
    if i % 200 == 0:
        print(i, loss)
# ⑥使用[[1, 11, 7, 9], [1, 3, 4, 3], [1, 1, 0, 1]]打印预测结果
y_test_predict = model(torch.Tensor([[1, 11, 7, 9], [1, 3, 4, 3], [1, 1, 0, 1]]))
print(torch.argmax(y_test_predict, 1))
# ⑦打印上面数据的准确率
y_train_predict = torch.argmax(model(x), 1)
acc = (y_train_predict == y.argmax(1)).float().mean().numpy()

# ⑧绘制损失函数图像
plt.plot(loss_list)
plt.show()
# ⑨打印模型准确率
print(acc)
# ⑩打印模型权重和截距
print(model.weight, model.bias)
